module mo_ps2str

  use shr_kind_mod, only : r8 => shr_kind_r8

  private

  public ps2str

contains

  subroutine ps2str(nw, zen, rsfc, tauu, omu, gu, dsdh, nid, radfld)
  !-----------------------------------------------------------------------------
  !   purpose:
  !   solve two-stream equations for multiple layers.  the subroutine is based
  !   on equations from:  toon et al., j.geophys.res., v94 (d13), nov 20, 1989.
  !   it contains 9 two-stream methods to choose from.  a pseudo-spherical
  !   correction has also been added.
  !-----------------------------------------------------------------------------
  !   parameters:
  !   nlevel  - integer, number of specified altitude levels in the working (i)
  !             grid
  !   zen     - real(r8), solar zenith angle (degrees)                          (i)
  !   rsfc    - real(r8), surface albedo at current wavelength                  (i)
  !   tauu    - real(r8), unscaled optical depth of each layer                  (i)
  !   omu     - real(r8), unscaled single scattering albedo of each layer       (i)
  !   gu      - real(r8), unscaled asymmetry parameter of each layer            (i)
  !   dsdh    - real(r8), slant path of direct beam through each layer crossed  (i)
  !             when travelling from the top of the atmosphere to layer i;
  !             dsdh(i,j), i = 0..nz-1, j = 1..nz-1
  !   nid     - integer, number of layers crossed by the direct beam when   (i)
  !             travelling from the top of the atmosphere to layer i;
  !             nid(i), i = 0..nz-1
  !   delta   - logical, switch to use delta-scaling                        (i)
  !             .true. -> apply delta-scaling
  !             .false.-> do not apply delta-scaling
  !   fdr     - real(r8), contribution of the direct component to the total     (o)
  !             actinic flux at each altitude level
  !   fup     - real(r8), contribution of the diffuse upwelling component to    (o)
  !             the total actinic flux at each altitude level
  !   fdn     - real(r8), contribution of the diffuse downwelling component to  (o)
  !             the total actinic flux at each altitude level
  !   edr     - real(r8), contribution of the direct component to the total     (o)
  !             spectral irradiance at each altitude level
  !   eup     - real(r8), contribution of the diffuse upwelling component to    (o)
  !             the total spectral irradiance at each altitude level
  !   edn     - real(r8), contribution of the diffuse downwelling component to  (o)
  !             the total spectral irradiance at each altitude level
  !-----------------------------------------------------------------------------

    use mo_params,    only : smallest, largest
    use mo_constants, only : d2r
    use ppgrid,       only : pver, pverp
    use mo_trislv,    only : tridec, trislv

    implicit none

    integer , intent(in)  :: nw
    integer , intent(in)  :: nid(0:pver)
    real(r8), intent(in)  :: zen
    real(r8), intent(in)  :: rsfc(nw)
    real(r8), intent(in)  :: tauu(pver,nw)
    real(r8), intent(in)  :: omu(pver,nw)
    real(r8), intent(in)  :: gu(pver,nw)
    real(r8), intent(in)  :: dsdh(0:pver,pver)
    real(r8), intent(out) :: radfld(pverp,nw)

    !-----------------------------------------------------------------------------
    ! 	... mu = cosine of solar zenith angle
    !           rsfc = surface albedo
    !           tauu =  unscaled optical depth of each layer
    !           omu  =  unscaled single scattering albedo
    !           gu   =  unscaled asymmetry factor
    !           klev = max dimension of number of layers in atmosphere
    !           nlayer = number of layers in the atmosphere
    !           nlevel = nlayer + 1 = number of levels
    !-----------------------------------------------------------------------------
    real(r8), parameter :: eps   = 1.e-3_r8
    real(r8), parameter :: pifs  = 1._r8
    real(r8), parameter :: fdn0  = 0._r8

    integer mrows
    integer pverm
    integer row
    integer lev
    integer i, ip1, wn
    integer j, jl, ju

    real(r8) precis, wrk
    real(r8) tempg
    real(r8) mu, suma
    real(r8) g, om
    real(r8) gam1, gam2, gam3, gam4
    real(r8), dimension(pver)    :: f, gi, omi
    real(r8), dimension(0:pverp) :: tauc, mu2
    real(r8), dimension(pver)    :: lam, taun, bgam
    real(r8), dimension(pver)    :: cdn
    real(r8), dimension(0:pverp,nw) :: tausla
    real(r8), dimension(pver,nw)  :: cup, cuptn, cdntn
    real(r8), dimension(pver,nw)  :: e1, e2, e3, e4
    real(r8), dimension(2*pver)    :: a, b, d, e
    real(r8), dimension(nw,2*pver) :: sub, main, super, y
!-----------------------------------------------------------------------------
! 	... for calculations of associated legendre polynomials for gama1,2,3,4
!           in delta-function, modified quadrature, hemispheric constant,
!           hybrid modified eddington-delta function metods, p633,table1.
!           w.e.meador and w.r.weaver, gas,1980,v37,p.630
!           w.j.wiscombe and g.w. grams, gas,1976,v33,p2440, 
!           uncomment the following two lines and the appropriate statements
!           further down.
!-----------------------------------------------------------------------------
    real(r8) :: expon, expon0, expon1, divisr, temp, up, dn
    real(r8) :: ssfc

    mrows = 2 * pver
    pverm = pver - 1

!-----------------------------------------------------------------------------
! 	... initial conditions:  pi*solar flux = 1;  diffuse incidence = 0
!-----------------------------------------------------------------------------
    precis = epsilon(precis)

    mu = cos(zen * d2r)
    wave_loop: do wn = 1, nw
!-----------------------------------------------------------------------------
!	... compute coefficients for each layer:
!           gam1 - gam4 = 2-stream coefficients, different for different approximations
!           expon0 = calculation of e when tau is zero
!           expon1 = calculation of e when tau is taun
!           cup and cdn = calculation when tau is zero
!           cuptn and cdntn = calc. when tau is taun
!           divisr = prevents division by zero
!-----------------------------------------------------------------------------
      tauc  (0:pverp)    = 0
      tausla(0:pverp,wn) = 0
      mu2   (0:pverp)    = sqrt(smallest)

!-----------------------------------------------------------------------------
! 	... delta-scaling. have to be done for delta-eddington approximation, 
!           delta discrete ordinate, practical improved flux method, delta function,
!           and hybrid modified eddington-delta function methods approximations
!-----------------------------------------------------------------------------
      f   (1:pver) = gu(:,wn) * gu(:,wn)
      gi  (1:pver) = (gu(:,wn) - f(1:pver)) / (1 - f(1:pver))
      omi (1:pver) = (1 - f(1:pver)) * omu(1:pver,wn) / (1 - omu(1:pver,wn) * f(1:pver))       
      taun(1:pver) = (1 - omu(1:pver,wn)*f(1:pver))*tauu(1:pver,wn)

!-----------------------------------------------------------------------------
! 	... calculate slant optical depth at the top of the atmosphere when zen>90.
!           in this case, higher altitude of the top layer is recommended which can 
!           be easily changed in gridz.f.
!-----------------------------------------------------------------------------
      if (zen > 90) then
        if (nid(0) < 0) then
          tausla(0,wn) = largest
        else
          ju = nid(0)
          tausla(0,wn) = 2 * dot_product(taun(1:ju),dsdh(0,1:ju))
        end if
      end if
      level_loop: do i = 1, pver
        g = gi(i)
        om = omi(i)
        tauc(i) = tauc(i-1) + taun(i)
!-----------------------------------------------------------------------------
! 	... stay away from 1 by precision.  for g, also stay away from -1
!-----------------------------------------------------------------------------
        tempg = min(abs(g), 1 - precis)
        g     = sign(tempg, g)
        om    = min(om, 1 - precis)
!-----------------------------------------------------------------------------
! 	... calculate slant optical depth
!-----------------------------------------------------------------------------
        if (nid(i) < 0) then
          tausla(i,wn) = largest
        else
          ju = min(nid(i), i)
          suma = dot_product(taun(1:ju), dsdh(i,1:ju))
          jl = min(nid(i), i) + 1
          tausla(i,wn) = suma + 2 * dot_product(taun(jl:nid(i)), dsdh(i,jl:nid(i)))
          if (tausla(i,wn) == tausla(i-1,wn)) then
            mu2(i) = sqrt(largest)
          else
            mu2(i) = (tauc(i) - tauc(i-1)) / (tausla(i,wn) - tausla(i-1,wn))
            mu2(i) = sign(max(abs(mu2(i)), sqrt(smallest)), mu2(i))
          end if
        end if
!-----------------------------------------------------------------------------
!	... the following gamma equations are from pg 16,289, table 1
!           eddington approximation(joseph et al., 1976, jas, 33, 2452):
!-----------------------------------------------------------------------------
        gam1 =  0.25_r8 * (7 - om * (4 + 3 * g))
        gam2 = -0.25_r8 * (1 - om * (4 - 3 * g))
        gam3 =  0.25_r8 * (2 - 3 * g * mu)
        gam4 = 1 - gam3
!-----------------------------------------------------------------------------
! 	... lambda = pg 16,290 equation 21
!           big gamma = pg 16,290 equation 22
!-----------------------------------------------------------------------------
        lam (i) = sqrt(gam1 * gam1 - gam2 * gam2)
        bgam(i) = (gam1 - lam(i)) / gam2
	      wrk = lam(i) * taun(i)
        if (wrk < 500) then
          expon = exp(-wrk)
        else
          expon = 0
        end if
!-----------------------------------------------------------------------------
! 	... e1 - e4 = pg 16,292 equation 44
!-----------------------------------------------------------------------------
        e1(i,wn) = 1 + bgam(i) * expon
        e2(i,wn) = 1 - bgam(i) * expon
        e3(i,wn) = bgam(i) + expon
        e4(i,wn) = bgam(i) - expon
!-----------------------------------------------------------------------------
! 	... the following sets up for the c equations 23, and 24
!           found on page 16,290
!           prevent division by zero (if lambda=1/mu, shift 1/mu^2 by eps = 1.e-3
!           which is approx equiv to shifting mu by 0.5*eps* (mu)**3
!-----------------------------------------------------------------------------
        if (tausla(i-1,wn) < 500) then
          expon0 = exp(-tausla(i-1,wn))
        else
          expon0 = 0
        end if
        if (tausla(i,wn) < 500) then
          expon1 = exp(-tausla(i,wn))
        else
          expon1 = 0
        end if
        divisr = lam(i) * lam(i) - 1.0_r8 / (mu2(i) * mu2(i))
        temp   = max(eps, abs(divisr))
        divisr = 1.0_r8 / sign(temp,divisr)
        up     = om * pifs * ((gam1 - 1.0_r8 / mu2(i)) * gam3 + gam4 * gam2) * divisr
        dn     = om * pifs * ((gam1 + 1.0_r8 / mu2(i)) * gam4 + gam2 * gam3) * divisr
!-----------------------------------------------------------------------------
! 	... cup and cdn are when tau is equal to zero
!           cuptn and cdntn are when tau is equal to taun
!-----------------------------------------------------------------------------
        cup  (i,wn) = up * expon0
        cdn  (i)    = dn * expon0
        cuptn(i,wn) = up * expon1
        cdntn(i,wn) = dn * expon1
      end do level_loop

!-----------------------------------------------------------------------------
!	... set up matrix
!           ssfc = pg 16,292 equation 37  where pi fs is one (unity).
!-----------------------------------------------------------------------------
	    if (tausla(pver,wn) < 500) then
        ssfc = rsfc(wn) * mu * exp(-tausla(pver,wn)) * pifs
      else
        ssfc = 0
      end if

!-----------------------------------------------------------------------------
! 	... the following are from pg 16,292  equations 39 - 43.
!           set up first row of matrix:
!-----------------------------------------------------------------------------
      a(1) = 0
      b(1) = e1(1,wn)
      d(1) = -e2(1,wn)
      e(1) = fdn0 - cdn(1)

!-----------------------------------------------------------------------------
! 	... set up odd rows 3 thru (mrows - 1):
!-----------------------------------------------------------------------------
      a(3:mrows-1:2) = e2(1:pverm,wn) * e3(1:pverm,wn) - e4(1:pverm,wn) * e1(1:pverm,wn)
      b(3:mrows-1:2) = e1(1:pverm,wn) * e1(2:pver ,wn) - e3(1:pverm,wn) * e3(2:pver ,wn)
      d(3:mrows-1:2) = e3(1:pverm,wn) * e4(2:pver ,wn) - e1(1:pverm,wn) * e2(2:pver ,wn)
      e(3:mrows-1:2) = e3(1:pverm,wn) * (cup(2:pver,wn) - cuptn(1:pverm,wn)) + e1(1:pverm,wn) * (cdntn(1:pverm,wn) - cdn(2:pver))

!-----------------------------------------------------------------------------
! 	... set up even rows 2 thru (mrows - 2): 
!-----------------------------------------------------------------------------
      a(2:mrows-2:2) = e2(2:pver ,wn) * e1(1:pverm,wn) - e3(1:pverm,wn) * e4(2:pver,wn)
      b(2:mrows-2:2) = e2(1:pverm,wn) * e2(2:pver ,wn) - e4(1:pverm,wn) * e4(2:pver,wn)
      d(2:mrows-2:2) = e1(2:pver ,wn) * e4(2:pver ,wn) - e2(2:pver ,wn) * e3(2:pver,wn)
      e(2:mrows-2:2) = (cup(2:pver,wn) - cuptn(1:pverm,wn)) * e2(2:pver,wn) - (cdn(2:pver) - cdntn(1:pverm,wn)) * e4(2:pver,wn)

!-----------------------------------------------------------------------------
! 	... set up last row of matrix at mrows:
!-----------------------------------------------------------------------------
      a(mrows) = e1(pver,wn) - rsfc(wn) * e3(pver,wn)
      b(mrows) = e2(pver,wn) - rsfc(wn) * e4(pver,wn)
      d(mrows) = 0
      e(mrows) = ssfc - cuptn(pver,wn) + rsfc(wn) * cdntn(pver,wn)

	    sub  (wn,1:mrows) = a(1:mrows)
	    main (wn,1:mrows) = b(1:mrows)
	    super(wn,1:mrows) = d(1:mrows)
	    y    (wn,1:mrows) = e(1:mrows)
    end do wave_loop

!-----------------------------------------------------------------------------
! 	... solve the system
!-----------------------------------------------------------------------------
    call tridec(nw, mrows, sub, main, super)
    call trislv(nw, mrows, sub, main, super, y)

!-----------------------------------------------------------------------------
!	... unfold solution of matrix, compute output fluxes
!-----------------------------------------------------------------------------
    do wn = 1, nw
!-----------------------------------------------------------------------------
! 	... the following equations are from pg 16,291  equations 31 & 32
!-----------------------------------------------------------------------------
      e(:mrows) = y(wn,:mrows)
      if (tausla(0,wn) < 500) then
        radfld(1,wn) = 2 * (fdn0 + e(1) * e3(1,wn) - e(2) * e4(1,wn) + cup(1,wn)) + exp(-tausla(0,wn))
      else
        radfld(1,wn) = 2 * (fdn0 + e(1) * e3(1,wn) - e(2) * e4(1,wn) + cup(1,wn))
      end if
      where (tausla(1:pver,wn) < 500)
        cdn(1:pver) = exp(-tausla(1:pver,wn))
      else where
        cdn(1:pver) = 0
      end where
      radfld(2:pverp,wn) = 2 * (e(1:mrows-1:2) * (e3(1:pver,wn) + e1(1:pver,wn)) &
                           + e(2:mrows:2) * (e4(1:pver,wn) + e2(1:pver,wn))      &
                           + cdntn(1:pver,wn) + cuptn(1:pver,wn)) + cdn(1:pver)
    end do

  end subroutine ps2str

end module mo_ps2str
